import numpy as np
from const import *
from newton1 import NewTon1
import matplotlib.pyplot as plt


class NewTon2:
    def __init__(self, A, b, x0=None):
        self.A = A
        self.AT = np.transpose(A)
        self.b = b
        self.n = A.shape[1]
        self.p = A.shape[0]
        if x0 is None:
            self.x0 = np.random.random(self.n)
        else:
            self.x0 = x0


    def f(self, x):
        return x @ np.log(x)

    def gradient(self, x):
        return np.log(x) + 1

    def hessian(self, x):
        res = np.zeros((len(x), len(x)))
        for i in range(len(x)):
            res[i][i] = 1.0 / x[i]
        return res

    def dy(self, x, v):
        mat = np.vstack((np.hstack((self.hessian(x),self.AT)),np.hstack((self.A,np.zeros((self.p,self.p))))))
        right = np.hstack((-1 * self.gradient(x) - self.AT @ v, - (self.A @ x - self.b)))
        dy = np.linalg.inv(mat) @ right
        return dy[0:self.n], dy[self.n:self.n + self.p]

    def norm(self, x, v):
        y = np.zeros((self.n + self.p))
        y[0:self.n] = self.gradient(x) + self.AT @ v
        y[self.n:self.n + self.p] = self.A @ x - self.b
        return np.linalg.norm(y)

    def ok(self, x):
        for i in x:
            if i <= 0:
                return False
        return True

    def search(self):
        v0 = np.random.random((self.p))
        x0 = self.x0
        ls_f = [self.f(x0)]
        while True:
            t = 1
            dx, dv = self.dy(x0, v0)
            nm = self.norm(x0, v0)
            if nm < epsilon:
                break
            x1 = x0 + t * dx
            v1 = v0 + t * dv
            while self.ok(x1) == False:
                t = t * beta
                x1 = x0 + t * dx
                v1 = v0 + t * dv
            while self.norm(x1, v1) > self.norm(x0, v0) * (1 - alpha * t):
                t = beta * t
                v1 = v0 + t * dv
                x1 = x0 + t * dx
            x0 = x1
            v0 = v1
            ls_f.append(self.f(x0))
            print(f"epoach {len(ls_f) - 1}  norm: {self.norm(x0, v0)}  f(x): {ls_f[-1]}")
        # n1 = NewTon1(self.A, x0)
        # x0, ls1 = n1.search()
        # ls_f.extend(ls1)
        return x0, v0, ls_f

    def search_plot(self):
        x, v, ls = self.search()
        print(f"x={x}")
        print(f"v={v}")
        print(f"f(x) = {ls[-1]}")
        plt.plot(range(len(ls)),ls)
        plt.xlabel("epoach")
        plt.ylabel("f(x)")
        plt.show()

if __name__ == "__main__":
    A = np.loadtxt("A.txt")
    b = np.loadtxt("b.txt")
    x0 = np.loadtxt("x0.txt")
    new = NewTon2(A, b, x0)
    new.search_plot()
